import random
import torchaudio
from torch.utils.data import Dataset, DataLoader


EMO2INX = {
    'neu':0,
    'ang':1,
    'hap':2,
    'sad':3
}


class AudioDataset(Dataset):
    def __init__(self, file_path, target_length_ms=5000, target_sample_rate=16000, norm_wav=True):
        self.file_path = file_path
        self.target_length_ms = target_length_ms
        self.norm_wav = norm_wav
        self.target_sample_rate = target_sample_rate
        self.data = []
        with open(file_path, 'r') as f:
            for line in f:
                eles = line.strip().split('|')
                if len(eles) == 2:
                    path = eles[0]
                    label = EMO2INX[eles[1]]
                    self.data.append((path, label))
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        wav_path, class_idx = self.data[idx]
        wavform, sample_rate = torchaudio.load(wav_path)
        if wavform.size(0) > 1:
            # mix to mono
            wavform = wavform.mean(dim=0, keepdim=True)
        if sample_rate != self.target_sample_rate:
            wavform = torchaudio.functional.resample(wavform, orig_freq=sample_rate, new_freq=self.target_sample_rate)
        wavform_length_ms = wavform.size(-1) * 1000 / self.target_sample_rate

        if wavform_length_ms < self.target_length_ms:
            num_repeats = int(self.target_length_ms / wavform_length_ms) + 1
            wavform = wavform.repeat(1, num_repeats)
        
        max_audio_start = wavform.size(-1) - int(self.target_length_ms* self.target_sample_rate / 1000)
        start_sample = random.randint(0, max_audio_start)
        end_sample = start_sample + int(self.target_length_ms * self.target_sample_rate / 1000)
        wavform = wavform[:, start_sample:end_sample]

        if self.norm_wav:
            # mono = wavform.mean(dim=1, keepdim=True)
            volume = wavform.pow(2).mean(dim=-1, keepdim=True).sqrt()
            scale = 1e-8 + volume
            wavform = wavform / scale
        
        return wavform.squeeze(0), class_idx

